import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from _data_process import cumulative_to_monthly
from _data_process import monthly_to_quarterly
from _data_process import calculate_yoy
from _data_plot import plot_double_axis_with_background
from _data_plot import plot_heatmap
from 关联分析数据 import merge_series_with_common_dates
from 关联分析数据 import add_same_trend_column


def plot_multi_line_chart(df, title="多城市地铁客运量与GDP季度同比", x_label="日期", y_label="同比增长率(%)"):
    """
    绘制多个城市的折线图，支持中文显示，GDP数据使用左轴，地铁客运量使用右轴
    
    参数:
        df: DataFrame，索引为日期，列为各个城市的地铁客运量季度同比数据，最后一列为GDP当季同比
        title: 图表标题，默认为"多城市地铁客运量与GDP季度同比"
        x_label: x轴标签，默认为"日期"
        y_label: y轴标签，默认为"同比增长率(%)"
    
    返回:
        fig: matplotlib的figure对象
    """
    # 设置中文字体支持
    plt.rcParams["font.family"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    
    # 创建图形和左轴
    fig, ax1 = plt.subplots(figsize=(14, 8))
    
    # 定义颜色序列，确保有足够的颜色
    colors = ['b', 'r', 'g', 'c', 'm', 'y', 'k', 'orange', 'purple', 'brown', 
              'pink', 'gray', 'olive', 'navy', 'teal', 'maroon', 'gold', 'lime']
    
    # 获取GDP数据（最后一列）
    gdp_column = df.columns[-1]
    metro_columns = df.columns[:-1]  # 除最后一列外的所有列
    
    # 绘制GDP数据在左轴（蓝色）
    ax1.plot(df.index, df[gdp_column], marker='o', linestyle='-', 
            color='blue', markersize=5, linewidth=2, label=gdp_column)
    ax1.set_xlabel(x_label, fontsize=12, labelpad=10)
    ax1.set_ylabel(f"{gdp_column} (%)", color='blue', fontsize=12)
    ax1.tick_params(axis='y', labelcolor='blue')
    
    # 创建右轴
    ax2 = ax1.twinx()
    
    # 为每个城市的地铁客运量数据绘制折线在右轴
    for i, column in enumerate(metro_columns):
        color = colors[i % len(colors)]
        ax2.plot(df.index, df[column], marker='o', linestyle='-', 
                color=color, markersize=4, linewidth=1.5, label=column)
    
    ax2.set_ylabel("地铁客运量季度同比 (%)", color='red', fontsize=12)
    ax2.tick_params(axis='y', labelcolor='red')
    
    # 设置标题
    ax1.set_title(title, fontsize=16, pad=20)
    
    # 美化x轴日期显示
    if hasattr(df.index, 'dtype') and 'datetime64' in str(df.index.dtype):
        ax1.xaxis.set_major_locator(mdates.AutoDateLocator())
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
        plt.xticks(rotation=45, ha='right')
    
    # 添加图例
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left', fontsize=10)
    
    # 添加网格线
    ax1.grid(True, linestyle='--', alpha=0.7)
    
    # 调整布局
    plt.tight_layout()
    
    return fig

df = pd.read_excel('data/20城市地铁客运量.xlsx', index_col=0)
df = monthly_to_quarterly(df)
ddf = df.copy(deep=True)
df = calculate_yoy(df)
gdp = pd.read_excel('data/GDP_不变价_当季同比.xlsx', index_col=0)

# 将GDP数据合并到地铁数据中，作为最后一列
df[gdp.columns[0]] = gdp.iloc[:, 0]
fg = plot_multi_line_chart(df)

df_sum = ddf.sum(axis=1)
df_sum = pd.DataFrame(df_sum)
sub_way_yoy = calculate_yoy(df_sum)


# 将两个时间序列合并为DataFrame，保留共同日期及最新两期非共同日期的数据
gdp_x_data = merge_series_with_common_dates(
    gdp.iloc[:,0],  # GDP数据的第一列
    sub_way_yoy.iloc[:,0],  
    "GDP当季同比",  # GDP列的命名
    "多城市地铁客运指数",  # x的命名
).iloc[5:,:]  # 去掉前5行数据（通常用于去除初期不稳定的数据）

    # 添加趋势分析列
    # 为DataFrame添加四列：第一列涨跌、第二列涨跌、是否同趋势以及index的月份
gdp_x_data = add_same_trend_column(gdp_x_data)

    # 绘制双轴对比图
    # 使用自定义函数绘制GDP与用电量的双轴趋势对比图，背景色表示同趋势情况
fig_gdp_x = plot_double_axis_with_background(
    gdp_x_data.iloc[:,:],  # 使用所有列的数据
    gdp_x_data.columns[0],  # 第一个列名（GDP当季同比）
    gdp_x_data.columns[1],  # 第二个列名（x_目标名）
    title=f"GDP当季同比与多城市地铁客运指数关联分析",  # 图表标题
)